Review: PR #745 — Expose row-group admission controls

Summary

This PR promotes the previously-internal row-group admission controls of AsyncTaskScheduler to a public RunConfig.row_group_admission policy (RowGroupAdmissionConfig + RowGroupAdmissionMode). It threads the public config through DatasetBuilder._prepare_async_run into the scheduler, applies the row-budget guard in both fixed and adaptive modes (previously adaptive-only), refactors a chunk of admission diagnostics emission, and adds matching architecture / Fern / docs/concepts/ documentation plus a plans/741/ design note.

The core API is small, well validated, and correctly preserves defaults: RunConfig() continues to behave like the prior fixed horizon of 3 row groups. Most of the volume is documentation and test coverage. The change is a clean, layered refactor — no import-direction violations and no leakage of engine internals into the config package.

Findings

Behavior change worth a release-note callout (medium)

_row_group_row_guard_allows previously short-circuited to True outside adaptive mode (async_scheduler.py old line: if not self._adaptive_row_group_admission: return True). It now applies unconditionally, with _max_admitted_rows derived as min(num_records, max(hard_cap * buffer_size, 8192)) when not explicit. The new test test_scheduler_derived_row_group_row_guard_blocks_extra_large_groups is parametrized over ["fixed", "adaptive"] and asserts the guard blocks 5,000-row groups under a derived guardrail of 8,192 — meaning fixed-mode runs that previously admitted oversized row groups will now be blocked.

In practice this is unlikely to bite typical users (buffer_size=1000, cap=3 gives a derived guard of 8,192 ≥ 3,000 worth of admitted rows, so the guard is a near no-op), but users running with a single row group materially larger than max(hard_cap * buffer_size, 8192) will see throttling they didn't before. The architecture doc describes the new guard, but neither the RunConfig.row_group_admission field docstring nor set_run_config mention that the default behavior in fixed mode tightened. Worth a one-line note in the architecture doc / changelog ("Fixed-mode runs now also enforce a derived max_admitted_rows guard; set max_admitted_rows explicitly to retain prior behavior or widen it").

_adaptive_row_group_block_reason lost its on-demand lookup (low)

Previously the function called _next_unadmitted_row_group_size() to scan self._row_groups for the next unadmitted group. Now the caller passes self._pending_row_group_admission_size, which is only set inside the _admit_row_groups loop while it's actively waiting for capacity (and cleared in finally). If _maybe_update_adaptive_row_group_target runs while _admit_row_groups is between iterations (after _dispatch_seeds, before reassignment), the value is stale-but-set; if _admit_row_groups has finished (_all_rgs_admitted) the early-return path covers it.

The cases are likely all benign because the actual admission loop independently re-evaluates the guard, but the new design ties an adaptive-target signal to the lifecycle of an unrelated coroutine variable. Two cheap mitigations: keep _next_unadmitted_row_group_size() (it's small), or document that _pending_row_group_admission_size is "current admission attempt, may be None between attempts" and accept no_pending_row_groups is a slightly noisier diagnostic. Not blocking.

max_admitted_rows or fallback (very minor)

async_scheduler.py:346:

self._max_admitted_rows = max_admitted_rows or self._max_admitted_rows_guardrail()

The constructor already validates max_admitted_rows >= 1, so or and is not None else are equivalent here. Prefer max_admitted_rows if max_admitted_rows is not None else self._max_admitted_rows_guardrail() for clarity — the or form silently falls through on a hypothetical 0 if the validator is ever loosened. Bikeshed.

Pydantic config — solid

• RowGroupAdmissionConfig cleanly enforces adaptive_initial_target is rejected in fixed mode and bounded in adaptive mode.
• RunConfig.validate_row_group_admission_budget correctly enforces max_admitted_rows >= buffer_size.
• Lazy-import wiring in data_designer.config.__init__ matches the existing pattern (entry in both TYPE_CHECKING block and _LAZY_NAME_TO_MOD_OBJ).
• _MAX_ROW_GROUP_ADMISSION_HORIZON = 64 and _MAX_EXPLICIT_ADMITTED_ROWS = 1_000_000 are reasonable safety caps; the values match the table in the docs.

Defaults preserve prior behavior — confirmed

• RunConfig().row_group_admission → mode=fixed, max_concurrent_row_groups=3, matching the prior scheduler default (test_run_config_exposes_default_row_group_admission_policy, test_prepare_async_run_threads_row_group_admission_config[default_fixed]).
• row_group_admission_source defaults to "default" in the scheduler signature; _prepare_async_run always passes "run_config" now, so the capacity-plan source flips from the old "dataset_builder" to "run_config" for production runs. Anything consuming that field on the public capacity plan should be aware. Worth checking if any telemetry consumer pins the literal "dataset_builder" for row-group-admission source.

Test coverage — strong

Six new config-level tests, one parametrized scheduler integration test (5 cases), two new async scheduler integration tests (fixed/adaptive parametrized), a default-row-guard scaling test, and an explicit-row-guard blocking test. Both fixed and adaptive paths are exercised end-to-end against MockSeedGenerator / SlowLLMBoundCellGenerator. The deletion of _next_unadmitted_row_group_size is implicitly covered by the adaptive-block-reason tests now passing next_size directly.

Style / structural compliance — clean

• from __future__ import annotations: present in modified files.
• Modern type syntax (int | None, list[str]).
• Absolute imports only.
• Import direction: dataset_builder.py correctly imports from data_designer.config.run_config; data_designer.interface.data_designer only updates a docstring. ✓
• The _emit_row_group_admission_event extraction is a nice cleanup — three callsites collapsed, less drift between the event name and the diagnostics reason.

Doc parity

architecture/dataset-builders.md, docs/concepts/architecture-and-performance.md, and fern/versions/latest/pages/concepts/architecture-and-performance.mdx are all updated with consistent wording. Code snippet imports data_designer.config as dd and from data_designer.interface import DataDesigner — both supported.

Verdict

Approve with minor suggestions. The shape of the public API matches existing patterns in the codebase, defaults are preserved, validation is tight, and tests cover both modes. The two things worth addressing before merge:

1. Note the fixed-mode row-guard tightening in the architecture doc / release notes (or relax the derived guard to remain a no-op for fixed mode unless explicitly set, if backward compatibility is intended).
2. Either restore _next_unadmitted_row_group_size() or add a one-line comment explaining the staleness contract of _pending_row_group_admission_size in _adaptive_row_group_block_reason.

Neither is a blocker.

Greptile Summary

This PR exposes RunConfig.max_concurrent_row_groups — previously an internal constant of 3 in AsyncTaskScheduler — as a validated public field, forwarding it through DatasetBuilder._prepare_async_run. The default is preserved, so existing behavior is unchanged.

• Adds a Pydantic field (ge=1, default 3) to RunConfig and passes it to AsyncTaskScheduler in dataset_builder.py.
• Adds unit tests for default, custom, and invalid (0) values; integration tests verify the kwarg is forwarded with the expected value.
• Updates the architecture doc and the Fern public docs with the new knob, memory implications, and checkpoint-timing caveats.

Important Files Changed

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant User
    participant RunConfig
    participant DatasetBuilder
    participant AsyncTaskScheduler

    User->>RunConfig: "RunConfig(max_concurrent_row_groups=N)"
    Note over RunConfig: Validates N >= 1
    User->>DatasetBuilder: builder.build(num_records, ...)
    DatasetBuilder->>DatasetBuilder: _prepare_async_run(...)
    DatasetBuilder->>AsyncTaskScheduler: "AsyncTaskScheduler(..., max_concurrent_row_groups=N, ...)"
    Note over AsyncTaskScheduler: _rg_semaphore = Semaphore(N)<br/>_row_group_admission_hard_cap = max(1, N)<br/>_row_group_admission_target = N (fixed mode)
    AsyncTaskScheduler-->>DatasetBuilder: scheduler
    DatasetBuilder-->>User: dataset

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant User
    participant RunConfig
    participant DatasetBuilder
    participant AsyncTaskScheduler

    User->>RunConfig: "RunConfig(max_concurrent_row_groups=N)"
    Note over RunConfig: Validates N >= 1
    User->>DatasetBuilder: builder.build(num_records, ...)
    DatasetBuilder->>DatasetBuilder: _prepare_async_run(...)
    DatasetBuilder->>AsyncTaskScheduler: "AsyncTaskScheduler(..., max_concurrent_row_groups=N, ...)"
    Note over AsyncTaskScheduler: _rg_semaphore = Semaphore(N)<br/>_row_group_admission_hard_cap = max(1, N)<br/>_row_group_admission_target = N (fixed mode)
    AsyncTaskScheduler-->>DatasetBuilder: scheduler
    DatasetBuilder-->>User: dataset

_{Reviews (6): Last reviewed commit: "Merge branch 'main' into etramel/fix/741..." | Re-trigger Greptile}

Nice work on this one, @eric-tramel — clean, well-contained change.

Summary

Exposes a single validated RunConfig.max_concurrent_row_groups (default 3, >= 1) and forwards it to the async scheduler's existing fixed row-group horizon, with matching docs and focused tests. The implementation matches the stated intent: the historical default is preserved end-to-end, and adaptive admission / completion machinery stay internal.

Findings

Suggestions — Take it or leave it

fern/versions/latest/pages/concepts/architecture-and-performance.mdx:106-116 — Concurrency-formula rewrite is broader than the PR title implies

• What: Beyond adding the new max_concurrent_row_groups section, the PR rewrites the "Concurrency Formula" block (buffer_size/remaining_cells_in_column → active_ready_model_cells/max_in_flight_tasks).
• Why: It's a genuine accuracy improvement — the new terms line up with the scheduler (per-model max_parallel_requests ceiling, max_in_flight_tasks as the task-lease cap) — so no objection. Just flagging that the formula correction rides along in a PR titled "make async row-group horizon configurable."
• Suggestion: A one-line note in the PR body ("also clarifies the concurrency formula") would help a future reader; no code change needed.

fern/.../architecture-and-performance.mdx:158-168 — Optionally tie the new section back to the formula

• What: The new section explains the memory/checkpoint tradeoff clearly but doesn't explicitly connect to the formula above (a wider horizon raises active_ready_model_cells).
• Suggestion: A sentence linking the two would help readers tuning throughput. Purely optional.

Existing bot comments (not this diff)

• The two open Greptile comments (P2 "stale pending size" and the _finalize_after_shutdown salvage question) both target async_scheduler.py, which isn't part of this changeset — they're left over from an earlier branch state. Worth resolving/dismissing so they don't muddy the PR, but nothing to fix here.

What Looks Good

• Default preserved end-to-end: RunConfig default 3 → scheduler max_concurrent_row_groups default 3 (previously implicit via the unset kwarg), so zero behavior change for existing users. Confirmed the builder at dataset_builder.py:956 is the only production instantiation.
• Consistent with siblings: field placement, docstring ordering, and the ge=1 validation all mirror max_in_flight_tasks; the new tests mirror the existing max_in_flight_tasks test trio (default / custom / reject-invalid).
• Tests cover the right seams: config default, custom value, and rejection of 0, plus a builder-wiring assertion that the value actually reaches the scheduler kwarg — the one most likely to catch a future regression.
• Architecture doc updated to keep architecture/dataset-builders.md truthful about the public surface.

Verdict

Ship it (with nits) — Only optional documentation-level suggestions. The code change is minimal, correct, well-tested, and preserves existing behavior.

Verified locally this pass: checked out the PR branch in a worktree and ran ruff check + ruff format --check on all four changed Python files — both pass cleanly. Also note the PR already has an approval from @andreatgretel.

This review was generated by an AI assistant.